"""
Celery任务监控和增强调度系统
提供任务状态监控、错误追踪、性能分析和自动恢复功能
"""
import time
import json
import asyncio
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from enum import Enum

from celery import Celery
from celery.result import AsyncResult
from celery.app.control import Control
from celery.utils.log import get_task_logger

from app.core.config import settings
from app.core.logging import get_logger

logger = get_logger(__name__)


class TaskStatus(Enum):
    """任务状态枚举"""
    PENDING = "PENDING"
    STARTED = "STARTED"
    SUCCESS = "SUCCESS"
    FAILURE = "FAILURE"
    RETRY = "RETRY"
    REVOKED = "REVOKED"


@dataclass
class TaskStats:
    """任务统计信息"""
    task_id: str
    name: str
    status: TaskStatus
    start_time: Optional[datetime]
    end_time: Optional[datetime]
    execution_time: Optional[float]
    result: Optional[Any]
    error: Optional[str]
    retries: int
    worker_name: Optional[str]


@dataclass
class WorkerStats:
    """工作者统计信息"""
    worker_name: str
    active_tasks: int
    processed_tasks: int
    failed_tasks: int
    avg_execution_time: float
    status: str
    last_heartbeat: Optional[datetime]


class CeleryMonitor:
    """Celery任务监控器"""
    
    def __init__(self, celery_app: Celery):
        self.celery_app = celery_app
        self.control = Control(celery_app)
        self.task_history: Dict[str, TaskStats] = {}
        self.worker_stats: Dict[str, WorkerStats] = {}
        
    async def get_task_status(self, task_id: str) -> Dict[str, Any]:
        """获取任务状态"""
        try:
            result = AsyncResult(task_id, app=self.celery_app)
            
            return {
                "task_id": task_id,
                "status": result.status,
                "result": result.result if result.successful() else None,
                "error": str(result.result) if result.failed() else None,
                "date_done": result.date_done.isoformat() if result.date_done else None,
                "traceback": result.traceback,
                "state": result.state
            }
        except Exception as e:
            logger.error(f"Failed to get task status for {task_id}: {e}")
            return {"error": str(e)}
    
    async def get_active_tasks(self) -> List[Dict[str, Any]]:
        """获取活动任务列表"""
        try:
            active_tasks = self.control.inspect().active()
            if not active_tasks:
                return []
            
            tasks = []
            for worker, worker_tasks in active_tasks.items():
                for task in worker_tasks:
                    tasks.append({
                        "worker": worker,
                        "task_id": task["id"],
                        "name": task["name"],
                        "args": task.get("args", []),
                        "kwargs": task.get("kwargs", {}),
                        "time_start": task.get("time_start"),
                        "acknowledged": task.get("acknowledged", False)
                    })
            
            return tasks
        except Exception as e:
            logger.error(f"Failed to get active tasks: {e}")
            return []
    
    async def get_scheduled_tasks(self) -> List[Dict[str, Any]]:
        """获取已调度任务列表"""
        try:
            scheduled_tasks = self.control.inspect().scheduled()
            if not scheduled_tasks:
                return []
            
            tasks = []
            for worker, worker_tasks in scheduled_tasks.items():
                for task in worker_tasks:
                    tasks.append({
                        "worker": worker,
                        "task_id": task["request"]["id"],
                        "name": task["request"]["name"],
                        "eta": task["request"].get("eta"),
                        "priority": task["request"].get("priority"),
                        "retries": task["request"].get("retries", 0)
                    })
            
            return tasks
        except Exception as e:
            logger.error(f"Failed to get scheduled tasks: {e}")
            return []
    
    async def get_worker_stats(self) -> Dict[str, Any]:
        """获取工作者统计信息"""
        try:
            stats = self.control.inspect().stats()
            if not stats:
                return {}
            
            worker_stats = {}
            for worker, worker_stat in stats.items():
                worker_stats[worker] = {
                    "active_tasks": worker_stat.get("total", {}).get("active", 0),
                    "processed_tasks": worker_stat.get("total", {}).get("task-received", 0),
                    "failed_tasks": worker_stat.get("total", {}).get("task-failed", 0),
                    "broker": worker_stat.get("broker", {}),
                    "pool": worker_stat.get("pool", {}),
                    "clock": worker_stat.get("clock")
                }
            
            return worker_stats
        except Exception as e:
            logger.error(f"Failed to get worker stats: {e}")
            return {}
    
    async def revoke_task(self, task_id: str, terminate: bool = False) -> Dict[str, Any]:
        """撤销任务"""
        try:
            result = self.control.revoke(task_id, terminate=terminate)
            return {
                "success": True,
                "task_id": task_id,
                "terminated": terminate,
                "result": result
            }
        except Exception as e:
            logger.error(f"Failed to revoke task {task_id}: {e}")
            return {"success": False, "error": str(e)}
    
    async def retry_task(self, task_id: str) -> Dict[str, Any]:
        """重试失败的任务"""
        try:
            # 获取任务结果
            result = AsyncResult(task_id, app=self.celery_app)
            
            if result.failed():
                # 重新发送任务
                new_task_id = self.celery_app.send_task(
                    result.name,
                    args=result.args,
                    kwargs=result.kwargs
                )
                
                return {
                    "success": True,
                    "original_task_id": task_id,
                    "new_task_id": new_task_id.id,
                    "message": "Task resubmitted for retry"
                }
            else:
                return {
                    "success": False,
                    "error": f"Task is not in failed state. Current state: {result.status}"
                }
                
        except Exception as e:
            logger.error(f"Failed to retry task {task_id}: {e}")
            return {"success": False, "error": str(e)}
    
    async def get_queue_stats(self) -> Dict[str, Any]:
        """获取队列统计信息"""
        try:
            # 获取Redis队列信息（假设使用Redis作为broker）
            import redis
            
            redis_client = redis.from_url(settings.CELERY_BROKER_URL)
            
            # 获取队列长度
            queues = ["celery", "celery_beat"]  # 可以根据需要扩展
            queue_stats = {}
            
            for queue in queues:
                queue_length = redis_client.llen(f"{queue}")
                queue_stats[queue] = {
                    "length": queue_length,
                    "name": queue
                }
            
            return queue_stats
        except Exception as e:
            logger.error(f"Failed to get queue stats: {e}")
            return {}
    
    async def cleanup_completed_tasks(self, older_than_hours: int = 24) -> Dict[str, Any]:
        """清理已完成的任务"""
        try:
            # 这里可以实现任务结果清理逻辑
            # 实际实现需要根据使用的结果后端来调整
            
            cleanup_time = datetime.now() - timedelta(hours=older_than_hours)
            cleaned_count = 0
            
            # 示例：清理任务历史记录
            task_ids_to_remove = []
            for task_id, task_stats in self.task_history.items():
                if task_stats.end_time and task_stats.end_time < cleanup_time:
                    task_ids_to_remove.append(task_id)
            
            for task_id in task_ids_to_remove:
                del self.task_history[task_id]
                cleaned_count += 1
            
            return {
                "success": True,
                "cleaned_count": cleaned_count,
                "older_than_hours": older_than_hours
            }
            
        except Exception as e:
            logger.error(f"Failed to cleanup tasks: {e}")
            return {"success": False, "error": str(e)}


class TaskScheduler:
    """增强的任务调度器"""
    
    def __init__(self, celery_app: Celery):
        self.celery_app = celery_app
        self.monitor = CeleryMonitor(celery_app)
        
    async def schedule_periodic_task(
        self,
        task_name: str,
        schedule: Dict[str, Any],
        args: Optional[List] = None,
        kwargs: Optional[Dict] = None,
        options: Optional[Dict] = None
    ) -> Dict[str, Any]:
        """调度周期性任务"""
        try:
            # 构建调度配置
            beat_schedule = {
                f"{task_name}_periodic": {
                    "task": task_name,
                    "schedule": schedule.get("interval", 300),  # 默认5分钟
                    "args": args or [],
                    "kwargs": kwargs or {},
                    "options": options or {}
                }
            }
            
            # 更新Celery配置
            self.celery_app.conf.beat_schedule.update(beat_schedule)
            
            return {
                "success": True,
                "task_name": task_name,
                "schedule": schedule,
                "message": "Periodic task scheduled successfully"
            }
            
        except Exception as e:
            logger.error(f"Failed to schedule periodic task {task_name}: {e}")
            return {"success": False, "error": str(e)}
    
    async def schedule_cron_task(
        self,
        task_name: str,
        cron_expression: str,
        args: Optional[List] = None,
        kwargs: Optional[Dict] = None
    ) -> Dict[str, Any]:
        """调度Cron任务"""
        try:
            from celery.schedules import crontab
            
            # 解析Cron表达式
            parts = cron_expression.split()
            if len(parts) != 5:
                raise ValueError("Invalid cron expression format")
            
            minute, hour, day_of_month, month, day_of_week = parts
            
            # 创建Crontab调度
            schedule = crontab(
                minute=minute,
                hour=hour,
                day_of_month=day_of_month,
                month_of_year=month,
                day_of_week=day_of_week
            )
            
            beat_schedule = {
                f"{task_name}_cron": {
                    "task": task_name,
                    "schedule": schedule,
                    "args": args or [],
                    "kwargs": kwargs or {}
                }
            }
            
            self.celery_app.conf.beat_schedule.update(beat_schedule)
            
            return {
                "success": True,
                "task_name": task_name,
                "cron_expression": cron_expression,
                "message": "Cron task scheduled successfully"
            }
            
        except Exception as e:
            logger.error(f"Failed to schedule cron task {task_name}: {e}")
            return {"success": False, "error": str(e)}
    
    async def get_scheduled_tasks_list(self) -> List[Dict[str, Any]]:
        """获取已调度的任务列表"""
        try:
            scheduled_tasks = []
            
            for task_name, task_config in self.celery_app.conf.beat_schedule.items():
                scheduled_tasks.append({
                    "name": task_name,
                    "task": task_config["task"],
                    "schedule": str(task_config["schedule"]),
                    "args": task_config.get("args", []),
                    "kwargs": task_config.get("kwargs", {}),
                    "options": task_config.get("options", {})
                })
            
            return scheduled_tasks
        except Exception as e:
            logger.error(f"Failed to get scheduled tasks list: {e}")
            return []
    
    async def remove_scheduled_task(self, task_name: str) -> Dict[str, Any]:
        """移除已调度的任务"""
        try:
            if task_name in self.celery_app.conf.beat_schedule:
                del self.celery_app.conf.beat_schedule[task_name]
                return {
                    "success": True,
                    "task_name": task_name,
                    "message": "Scheduled task removed successfully"
                }
            else:
                return {
                    "success": False,
                    "error": f"Task {task_name} not found in schedule"
                }
        except Exception as e:
            logger.error(f"Failed to remove scheduled task {task_name}: {e}")
            return {"success": False, "error": str(e)}


# 增强的Celery应用配置
def create_enhanced_celery_app(app_name: str = "celery_app") -> Celery:
    """创建增强的Celery应用"""
    celery_app = Celery(app_name)
    
    # 基础配置
    celery_app.conf.update(
        # Broker配置
        broker_url=settings.CELERY_BROKER_URL or "redis://localhost:6379/1",
        result_backend=settings.CELERY_RESULT_BACKEND or "redis://localhost:6379/2",
        
        # 序列化配置
        task_serializer="json",
        accept_content=["json"],
        result_serializer="json",
        
        # 时区配置
        timezone="Asia/Shanghai",
        enable_utc=True,
        
        # 工作者配置
        worker_prefetch_multiplier=1,
        task_acks_late=True,
        worker_max_tasks_per_child=1000,
        
        # 性能优化配置
        worker_concurrency=4,
        task_time_limit=300,  # 5分钟超时
        task_soft_time_limit=240,  # 4分钟软超时
        
        # 重试配置
        task_acks_on_failure_or_timeout=False,
        task_reject_on_worker_lost=True,
        
        # 监控配置
        worker_send_task_events=True,
        task_send_sent_event=True,
        
        # Beat调度器配置
        beat_schedule={},
        beat_scheduler="celery.beat:PersistentScheduler",
        beat_schedule_filename="celerybeat-schedule.db"
    )
    
    return celery_app


# 创建全局实例
from app.tasks.celery_app import celery_app

# 创建监控和调度器实例
celery_monitor = CeleryMonitor(celery_app)
task_scheduler = TaskScheduler(celery_app)